import cv2
import numpy as np
import onnxruntime
import os
def process_depth(image_path, onnx_path):
    # 初始化ONNX模型
    session = onnxruntime.InferenceSession(onnx_path)
    
    # 预处理
    def preprocess(img_path):
        image = cv2.cvtColor(cv2.imread(img_path), cv2.COLOR_BGR2RGB)
        image = cv2.resize(image, (518, 518), interpolation=cv2.INTER_CUBIC)
        image = image.astype(np.float32) / 255.0
        image = (image - [0.485, 0.456, 0.406]) / [0.229, 0.224, 0.225]
        return image.transpose(2, 0, 1)[np.newaxis, ...].astype(np.float32)

    # 执行推理
    input_tensor = preprocess(image_path)
    depth_output = session.run(
        [session.get_outputs()[0].name],
        {session.get_inputs()[0].name: input_tensor}
    )[0].squeeze()

    # 后处理
    original = cv2.imread(image_path)
    h, w = original.shape[:2]
    resized_depth = cv2.resize(depth_output, (w, h), interpolation=cv2.INTER_LINEAR)
    filename = os.path.basename(image_path)  # 新增：提取文件名
    # print(resized_depth[0,500,500])
    # print(max(resized_depth))
    # print(min(resized_depth))
    # 保存结果
    normalized_depth = cv2.normalize(resized_depth, None, 0, 255, cv2.NORM_MINMAX, cv2.CV_8U)
    # cv2.imwrite(f'./static/Saveload/depth/black_{filename}', normalized_depth)
    # cv2.imwrite(f'./static/Saveload/depth/color_{filename}', cv2.applyColorMap(normalized_depth, cv2.COLORMAP_INFERNO))
    return resized_depth
if __name__ == "__main__":
    process_depth(
        image_path=r'E:\AIproject\project1\fastApiProject\static\SaveLoad\uploadImages\20150921_131234_image11.png',
        onnx_path=r'E:\AIproject\project1\fastApiProject\depth_anything_v2_vitb.onnx'
    )
 